#include "resolv.h"
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <stdio.h>
#include <ctype.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include "udns/udns.h"
#include "io.h"
#include "misc.h"
#include "log.h"


struct resolv {
    struct resolv_ctx *ctx;
    const char *domain; /* be careful we use strdup */
    struct dns_query *queris[2]; /* [ QUERY_A, QUERY_AAAA ] */
    struct addr *addr; /* resolve result */

    /* maybe more than one request same domain name, we merge it */
    size_t reqstn;
    struct {
        resolv_cb *cb;
        void *udata;
    } *reqsts;
};

#define QUERY_IS_FINISH(resolv) \
    (!(resolv)->queris[QUERY_A] && !(resolv)->queris[QUERY_AAAA])
#define QUERY_IS_INVALID QUERY_IS_FINISH


static int is_ip(const char *domain, struct addr *addr) {

    if (inet_pton(AF_INET, domain, addr->addr) == 1)
        addr->addr->family = AF_INET;
    else if (inet_pton(AF_INET6, domain, addr->addr) == 1)
        addr->addr->family = AF_INET6;
    else
        return 0;

    addr->cnt = 1;
    addr->expir = -1;

    return 1;
}

static int hosts_cache_one(struct dict *dict, const char *name, struct addr *addr) {
    struct addr *oldaddr = (struct addr *)dict_get(dict, name, NULL);
    int oldaddrn = oldaddr ? oldaddr->cnt : 0;
    size_t namelen = strlen(name);
    /* we stored the key(name) and value(addr) into single memory chunk */
    void *mem = malloc(namelen + 1 /* key */
            + sizeof(*oldaddr) + sizeof(oldaddr->addr[0]) * oldaddrn); /* value */
    if (mem) {
        struct addr *newaddr = mem + namelen + 1;
        /* key preparation */
        memcpy(mem, name, namelen + 1);
        /* value preparation */
        if (oldaddrn) {
            /* copy old address */
            memcpy(newaddr, oldaddr, sizeof(*oldaddr) + sizeof(oldaddr->addr[0]) * (oldaddrn - 1));
            /* copy additional address */
            newaddr->addr[oldaddrn] = addr->addr[0];
            newaddr->cnt += 1;
        } else {
            *newaddr = *addr;
        }
        /* insert */
        const char *old;
        if (dict_set(dict, mem, newaddr, &old, NULL) < 0)
            free(mem);
        else {
            free((void *)old);
            return 0;
        }
    }

    errno = ENOMEM;
    return -1;
}

static void load_hosts(struct resolv_ctx *ctx) {
    const char *sep = " \t\r\n";
    char buf[512];
    size_t left = 0;
    FILE *fp;
    char *line, *brkt, *tmp;

    struct addr addr;
    int lineflags = 0;

#define LINE_WAS_TERM       1
#define LINE_TO_BE_CONTINUE 2
#define LINE_IS_COMMENT     4 /* next is comment data */
#define MARK_NEXT_LINE_COMMENT(flags) \
        do { \
            if (!(flags & LINE_WAS_TERM)) { \
                flags |= LINE_IS_COMMENT; \
            } \
        } while (0)

    if (!(fp = fopen("/etc/hosts", "r")))
        return;

get_next_line:
    while ((line = fgets(buf + left, sizeof(buf) - left, fp))) {
        left = 0;
        int c = line[strlen(line) - 1];
        if (c == '\r' || c == '\n') {
            lineflags |= LINE_WAS_TERM;
        } else {
            lineflags &= ~LINE_WAS_TERM;
        }
        if (lineflags & LINE_IS_COMMENT) {
            if (lineflags & LINE_WAS_TERM)
                lineflags &= ~LINE_IS_COMMENT;
            continue;
        }

#define CHECK_STR_AND_CONTINUE(ptr) if (!ptr) continue; \
        if (*ptr == '#') { MARK_NEXT_LINE_COMMENT(lineflags); continue; }

        if (!(lineflags & LINE_TO_BE_CONTINUE)) {
            char *ip = strtok_r(buf, sep, &brkt);
            CHECK_STR_AND_CONTINUE(ip);
            if (!is_ip(ip, &addr))
                continue;
            tmp = NULL;
        } else {
            tmp = buf;
        }

        /* the line what at buf maybe not complete,
         * we should do that correctly */
        char *name, *name_tmp;
        name = strtok_r(tmp, sep, &brkt);
        CHECK_STR_AND_CONTINUE(name);
        while ((name_tmp = strtok_r(NULL, sep, &brkt))) {
            hosts_cache_one(ctx->hosts, name, &addr);
            name = name_tmp;
            if (*name == '#') {
                MARK_NEXT_LINE_COMMENT(lineflags);
                goto get_next_line;
            }
        }

        if (lineflags & LINE_WAS_TERM) {
            hosts_cache_one(ctx->hosts, name, &addr);
        } else {
            left = strlen(name);
            if (left >= sizeof(buf) - 1) {
                /* invalid data */
                MARK_NEXT_LINE_COMMENT(lineflags);
                left = 0;
            }
            memmove(buf, name, left);
        }
    }
}

static void resolv_timeout_cb(EV_P_ ev_timer *watcher, int revents)
{
    if (revents & EV_TIMER) {
        dns_timeouts(watcher->data, 30, ev_now(EV_A));
    }
}

static void resolv_event_handler(EV_P_ ev_io *w, int revents) {
    if (revents & EV_READ)
        dns_ioevent(w->data, ev_now(EV_A));
}

static void dns_timer_setup_cb(struct dns_ctx *ctx, int timeout, void *data)
{
    struct resolv_ctx *rctx = data;

    ev_timer_stop(rctx->evloop, &rctx->evtimer);

    if (ctx && timeout >= 0) {
        ev_timer_set(&rctx->evtimer, timeout, 0.);
        ev_timer_start(rctx->evloop, &rctx->evtimer);
    }
}

struct resolv_ctx *resolv_ctx_create(time_t dns_cache_expire, struct ev_loop *evloop) {
    struct resolv_ctx *ctx;
    struct dns_ctx *dnsctx = &dns_defctx;

    if (!(ctx = malloc(sizeof(*ctx))) ||
            !(ctx->hosts = dict_create()) ||
            !(ctx->cache.dict = dict_create()) ||
            !(ctx->queryings = dict_create())) {
        LOGE("%s() failed: Not enough memory", __FUNCTION__);
        dict_destroy(ctx->cache.dict, free, free);
        dict_destroy(ctx->hosts, free, free);
        free(ctx);
        return NULL;
    }

    dns_init(dnsctx, 0);
    ctx->udns_ctx = dnsctx;
    if (!dns_nsrv(dnsctx)) {
        dns_add_serv(dnsctx, "119.29.29.29");
        dns_add_serv(dnsctx, "114.114.114.114");
    }

    if ((ctx->query_sock = dns_open(ctx->udns_ctx)) < 0) {
        free(ctx);
        LOGE("%s() failed: %s", __FUNCTION__, strerror(errno));
        return NULL;
    }
    nonblocking(ctx->query_sock);

    ctx->cache.expire = dns_cache_expire;

    load_hosts(ctx);

    ctx->evloop = evloop;
    ev_timer_init(&ctx->evtimer, resolv_timeout_cb, 0., 0.);
    ctx->evtimer.data = dnsctx;
    ev_io_init(&ctx->evio, resolv_event_handler, ctx->query_sock, EV_READ);
    ctx->evio.data = dnsctx;
    ev_io_start(evloop, &ctx->evio);

    dns_set_tmcbck(ctx->udns_ctx, dns_timer_setup_cb, ctx);

    return ctx;
}

void resolv_ctx_destroy(struct resolv_ctx *ctx) {
    dns_free(ctx->udns_ctx);
    dict_destroy(ctx->hosts, free, NULL);
    dict_destroy(ctx->cache.dict, free, free);
    free(ctx);
}

static void check_is_finish_and_callback(struct resolv *resolv) {
    if (QUERY_IS_FINISH(resolv)) {
        if (!resolv->addr)
            LOGI("Resolve \"%s\" failed: %s",
                    resolv->domain, "domain name not provided, or not known");

        /* callback */
        size_t i;
        for (i = 0; i < resolv->reqstn; ++i)
            if (resolv->reqsts[i].cb)
                resolv->reqsts[i].cb(resolv->domain, resolv->addr, resolv->reqsts[i].udata);

        /* Caching */
        /* if doamin name not found, we save a empty address */
        if (!resolv->addr && (resolv->addr = malloc(sizeof(*resolv->addr))))
            resolv->addr->cnt = 0;
        int cache_failed = 1;
        if (resolv->addr) {
            ((struct addr *)resolv->addr)->expir = resolv->addr->cnt ? time(NULL) + resolv->ctx->cache.expire : 10;
            const char *oldkey;
            const void *oldvalue;
            if (dict_set(resolv->ctx->cache.dict, resolv->domain, resolv->addr, &oldkey, &oldvalue) < 0) {
                /* caching failed, clean result */
                free(resolv->addr);
            } else {
                /* caching success, clean old memory */
                free((void *)oldkey);
                free((void *)oldvalue);
                cache_failed = 0;
            }
        }

        /* clean */
        dict_del(resolv->ctx->queryings, resolv->domain, NULL, NULL);
        if (cache_failed)
            free((void *)resolv->domain);
        free(resolv->reqsts);
        free(resolv);
    }
}

#define dns_rr_parse(resolv, result, ver, af) \
    do { \
        if ((result) && (result)->dnsa##ver##_nrr > 0) { \
            struct addr *new, *old = (struct addr *)(resolv)->addr; \
            int oldcnt = (old) ? (old)->cnt : 0; \
            int newcnt = (result)->dnsa##ver##_nrr; \
            new = realloc(old, sizeof(*new) + \
                    sizeof(new->addr[0]) * (oldcnt + newcnt - 1)); \
            if (new) { \
                int i; \
                for (i = 0; i < (result)->dnsa##ver##_nrr; ++i) { \
                    new->addr[oldcnt + i].in##ver = (result)->dnsa##ver##_addr[i]; \
                    new->addr[oldcnt + i].family = af; \
                } \
                new->cnt = oldcnt + newcnt; \
                (resolv)->addr = new; \
            } \
        } \
    } while (0)

static void dns_query_a4_cb(struct dns_ctx *ctx, struct dns_rr_a4 *result, void *data) {
    struct resolv *resolv = data;
    dns_rr_parse(resolv, result, 4, AF_INET);
    free(result);
    resolv->queris[QUERY_A] = NULL;
    check_is_finish_and_callback(resolv);
}

static void dns_query_a6_cb(struct dns_ctx *ctx, struct dns_rr_a6 *result, void *data) {
    struct resolv *resolv = data;
    dns_rr_parse(resolv, result, 6, AF_INET6);
    free(result);
    resolv->queris[QUERY_AAAA] = NULL;
    check_is_finish_and_callback(resolv);
}

const struct addr *in_hosts_or_in_cache(struct resolv_ctx *ctx, const char *name) {
    const struct addr *addr;
    if (((addr = dict_get(ctx->hosts, name, NULL)) ||
        (addr = dict_get(ctx->cache.dict, name, NULL))) &&
            (addr->expir > time(NULL) || addr->expir < 0)) {
        return addr;
    }
    return NULL;
}

struct resolv *is_querying(struct resolv_ctx *ctx, const char *domain,
        resolv_cb *cb, void *udata) {
    struct resolv *resolv;
    const char *key;

    if ((resolv = (struct resolv *)dict_get(ctx->queryings, domain, &key))) {
        resolv->reqsts = realloc(resolv->reqsts, sizeof(*resolv->reqsts) * (resolv->reqstn + 1));
        if (resolv->reqsts) {
            resolv->reqsts[resolv->reqstn].cb = cb;
            resolv->reqsts[resolv->reqstn].udata = udata;
            ++resolv->reqstn;
        }
    }

    return resolv;
}

/* if "domain" is IP string or in /etc/hosts or was caching,
 * immediately callback the "cb", and return NULL with errno set to 0.
 *
 * on success add query, return the resolv memory pointer.
 *
 * otherwise return NULL with errno set to ENOMEM or EINVAL.
 *
 * Careful: on the "cb" routine, do not touch the "resolv" (return of this
 *  function), its memory was freed.
 */
struct resolv *resolv_query(struct resolv_ctx *ctx, const char *domain,
        resolv_cb *cb, void *udata) {

    struct addr addr;
    const struct addr *addrp;
    struct resolv *resolv;

    if (is_ip(domain, &addr)) {
        cb(domain, &addr, udata);
        errno = 0; // mark success
        return NULL;
    }

    if ((addrp = in_hosts_or_in_cache(ctx, domain))) {
        cb(domain, addrp, udata);
        errno = 0; // mark success
        return NULL;
    }

    if ((resolv = is_querying(ctx, domain, cb, udata)))
        return resolv;

    if (!(resolv = malloc(sizeof(*resolv))) ||
            !(resolv->reqsts = malloc(sizeof(*resolv->reqsts))) ||
            !(domain = strdup(domain))) {
        free(resolv->reqsts);
        free(resolv);
        errno = ENOMEM;
        return NULL;
    }

    resolv->ctx = ctx;
    resolv->domain = domain;
    resolv->addr = NULL;
    resolv->reqstn = 1;
    resolv->reqsts[0].cb = cb;
    resolv->reqsts[0].udata = udata;

    resolv->queris[QUERY_A] = dns_submit_a4(ctx->udns_ctx, domain, 0, dns_query_a4_cb, resolv);
    resolv->queris[QUERY_AAAA] = dns_submit_a6(ctx->udns_ctx, domain, 0, dns_query_a6_cb, resolv);

    if (QUERY_IS_INVALID(resolv)) {
        int dnserr = dns_status(ctx->udns_ctx);
        LOGW("Failed to submit DNS query: %s", dns_strerror(dnserr));
        if (dnserr == DNS_E_NOMEM)
            errno = ENOMEM;
        else
            errno = EINVAL;
        free((void *)resolv->domain);
        free(resolv);
        return NULL;
    }

    if (dict_set(ctx->queryings, domain, resolv, NULL, NULL) < 0) {
        resolv_cancel(resolv, cb, udata);
        free((void *)resolv->domain);
        free(resolv);
        errno = ENOMEM;
        return NULL;
    }

    return resolv;
}

void resolv_cancel(struct resolv *resolv, resolv_cb *cb, void *udata) {
    size_t i;

    if (!resolv || !cb)
        return;

#define REAL_CANCEL 0

    for (i = 0; i < resolv->reqstn; ++i)
        if (resolv->reqsts[i].udata == udata && resolv->reqsts[i].cb == cb) {
#if REAL_CANCEL
            if (i != resolv->reqstn - 1)
                memcpy(resolv->reqsts + i, resolv->reqsts + i + 1,
                        sizeof(resolv->reqsts) * (resolv->reqstn - i - 1));
             --resolv->reqstn;
#else
            /* we just mark it */
             resolv->reqsts[i].cb = NULL;
#endif
            break;
        }

#if REAL_CANCEL
    if (resolv->reqstn)
        return;

    /* now, we need really cancel */
    for (i = 0; i < ARRAY_LEN(resolv->queris); ++i)
        if (resolv->queris[i]) {
            dns_cancel(resolv->ctx->udns_ctx, resolv->queris[i]);
            free(resolv->queris[i]);
            resolv->queris[i] = NULL;
        }

    /* maybe some part (QUERY_A or QUERY_AAAA) completed, we need clean it */
    if (resolv->addr) {
        free((void *)resolv->addr);
        resolv->addr = NULL;
    }

    dict_del(resolv->ctx->queryings, resolv->domain, NULL, NULL);
    free((void *)resolv->domain);
    free(resolv->reqsts);
    free(resolv);
#endif
}
